# #!/usr/bin/env python
# # -*- coding: utf-8 -*-
# # @Time    : 2025/7/29
# # @Author  : YunZhen
# # @File    : processor.py
# # @Software: PyCharm
"""

"""
import math
from loguru import logger
from typing import Iterable, Iterator, Tuple, List, Dict, Any, Optional, TypeVar, Generic, Callable

T = TypeVar('T')  # 通用类型变量


class BaseBatchProcessor(Generic[T]):
    """
    批处理器基类 - 提供通用批处理功能
    提供：
    1. 元素总数
    2. 总批次数量
    3. 当前批次索引
    4. 当前批次的起始和结束索引
    5. 当前批次数据
    子类需要实现 __next__ 方法
    """

    def __init__(
            self,
            iterable: Iterable[T],
            batch_size: int = 100,
            total_elements: Optional[int] = None
    ):
        """
        初始化批次处理器

        :param iterable: 要处理的可迭代对象
        :param batch_size: 每批的大小，默认为100
        :param total_elements: 元素总数（可选），如果未提供将尝试获取
        """
        self.iterable = iterable
        self.batch_size = batch_size

        # 确定元素总数
        if total_elements is not None:
            self.total_elements = total_elements
        else:
            try:
                self.total_elements = len(iterable)  # type: ignore
            except TypeError:
                # 对于没有长度信息的迭代器(生成器、文件对象）不支持 len()，先转换为list
                self.iterable = list(iterable)
                self.total_elements = len(self.iterable)

        # 计算总批次
        self.total_batches = math.ceil(self.total_elements / batch_size)

        # 初始化状态
        self.current_batch = 0
        self._iterator = iter(self.iterable)

    def __iter__(self) -> Iterator[Any]:
        """迭代器接口"""
        self.current_batch = 0
        self._iterator = iter(self.iterable)  # 重置迭代器
        return self

    def __next__(self) -> Any:
        """获取下一批数据 - 子类必须实现"""
        raise NotImplementedError("子类必须实现 __next__ 方法")

    def get_batch_info(self) -> dict:
        """获取当前批次信息"""
        return {
            "total_elements": self.total_elements,
            "total_batches": self.total_batches,
            "current_batch": self.current_batch,
            "batch_size": self.batch_size,
            "current_start": (self.current_batch - 1) * self.batch_size if self.current_batch > 0 else 0,
            "current_end": min(self.current_batch * self.batch_size, self.total_elements)
        }

    def reset(self) -> None:
        """重置处理器状态"""
        self.current_batch = 0
        self._iterator = iter(self.iterable)


class ListBatchProcessor(BaseBatchProcessor[T]):
    """
    列表批处理器 - 高效处理列表类型数据的分批-使用索引切片代替迭代器遍历
    """

    def __init__(
            self,
            iterable: Iterable[T],
            batch_size: int = 100,
            total_elements: Optional[int] = None
    ):
        """
        初始化列表批处理器
        特别注意：对于不支持索引的可迭代对象（如生成器），会先转换为列表
        """
        super().__init__(iterable, batch_size, total_elements)

        # 确保可索引：如果原始数据不是列表，转换为列表
        if not isinstance(self.iterable, list):
            self.iterable = list(self.iterable)

    def __next__(self) -> Tuple[List[T], int, int]:
        """获取下一批数据 - 使用索引切片"""
        if self.current_batch >= self.total_batches:
            raise StopIteration

        # 计算当前批次的索引范围
        start_idx = self.current_batch * self.batch_size
        end_idx = min((self.current_batch + 1) * self.batch_size, self.total_elements)

        # 使用切片获取批次数据
        batch_data = self.iterable[start_idx:end_idx]  # 关键优化点

        # 更新批次计数
        self.current_batch += 1

        return batch_data, start_idx, end_idx


class DictBatchProcessor(BaseBatchProcessor[Dict[str, Any]]):
    """
    字典批处理器 - 处理字典类型数据的分批

    提供额外的字典处理功能：
    1. 按字段过滤
    2. 按字典键过滤
    """

    def __init__(
            self,
            iterable: Iterable[Dict[str, Any]],
            batch_size: int = 100,
            total_elements: Optional[int] = None,
            excls_filter: Optional[str] = None,
            filter_func: Optional[callable] = None
    ):
        """
        初始化字典批处理器

        :param excls_filter: 过滤的键字段
        :param filter_func: 键过滤函数，返回True保留
        """
        # 保存原始数据
        self.original_data = iterable
        self.key_field = excls_filter
        self.key_filter = filter_func

        # 应用过滤条件
        filtered_data = self._apply_filters()

        # 使用过滤后的数据调用父类初始化
        super().__init__(filtered_data, batch_size, total_elements)

    def _apply_filters(self) -> List[Dict[str, Any]]:
        """
        应用过滤条件并返回符合条件的字典列表
        """
        filtered_data = []

        for item in self.original_data:
            # 应用键字段过滤
            if self.key_field and self.key_field not in item:
                logger.debug(f"跳过缺失键字段 '{self.key_field}' 的项目")
                continue

            # 应用键过滤函数
            if self.key_filter:
                if not self.key_filter(item):
                    logger.debug(f"跳过不符合键过滤条件的项目")
                    continue

            filtered_data.append(item)

        return filtered_data

    def __next__(self) -> Tuple[List[Dict[str, Any]], int, int]:
        """获取下一批数据- 使用迭代器实现"""
        if self.current_batch >= self.total_batches:
            raise StopIteration

        # 计算当前批次的索引范围
        start_idx = self.current_batch * self.batch_size
        end_idx = min((self.current_batch + 1) * self.batch_size, self.total_elements)

        # 收集当前批次的数据
        batch_data = []
        items_to_fetch = min(self.batch_size, self.total_elements - start_idx)

        for _ in range(items_to_fetch):
            try:
                batch_data.append(next(self._iterator))
            except StopIteration:
                # 防止数据不一致
                if len(batch_data) == 0:
                    raise
                break

        # 更新批次计数
        self.current_batch += 1

        return batch_data, start_idx, end_idx

    def get_filtered_count(self) -> int:
        """获取过滤后的元素数量"""
        return self.total_elements

    def get_original_count(self) -> int:
        """获取原始元素数量"""
        try:
            return len(self.original_data)
        except TypeError:
            return len(list(self.original_data))

    def apply_new_filters(
            self,
            key_field: Optional[str] = None,
            key_filter: Optional[callable] = None
    ):
        """应用新的过滤条件并重置处理器"""
        self.key_field = key_field
        self.key_filter = key_filter

        # 重新应用过滤
        filtered_data = self._apply_filters()

        # 更新父类属性
        self.iterable = filtered_data
        self.total_elements = len(filtered_data)
        self.total_batches = math.ceil(self.total_elements / self.batch_size)
        self.reset()


class SmartBatchProcessor(BaseBatchProcessor[T]):
    """
    智能批处理器 - 添加进度回调功能
    """

    def __init__(
            self,
            iterable: Iterable[T],
            batch_size: int = 100,
            total_elements: Optional[int] = None,
            callback: Optional[Callable[[int, int], None]] = None
    ):
        """
        初始化智能批次处理器

        :param callback: 进度回调函数，格式为 callback(processed_count, total)
        """
        super().__init__(iterable, batch_size, total_elements)
        self.callback = callback
        self.processed_count = 0

    def __next__(self) -> Tuple[List[T], int, int]:
        """获取下一批数据"""
        batch_data, start_idx, end_idx = super().__next__()

        # 更新处理计数
        self.processed_count += len(batch_data)

        # 调用进度回调
        if self.callback:
            self.callback(self.processed_count, self.total_elements)

        return batch_data, start_idx, end_idx

    def reset(self) -> None:
        """重置处理器状态"""
        super().reset()
        self.processed_count = 0


class SmartDictBatchProcessor(DictBatchProcessor):
    """
    智能字典批处理器 - 结合字典处理和进度回调
    """

    def __init__(
            self,
            iterable: Iterable[Dict[str, Any]],
            batch_size: int = 100,
            total_elements: Optional[int] = None,
            callback: Optional[Callable[[int, int], None]] = None,
            key_field: Optional[str] = None,
            value_range: Optional[Tuple[Any, Any]] = None
    ):
        """
        初始化智能字典批处理器

        :param callback: 进度回调函数
        """
        super().__init__(iterable, batch_size, total_elements, key_field, value_range)
        self.callback = callback
        self.processed_count = 0

    def __next__(self) -> Tuple[List[Dict[str, Any]], int, int]:
        """获取下一批数据"""
        batch_data, start_idx, end_idx = super().__next__()

        # 更新处理计数
        self.processed_count += len(batch_data)

        # 调用进度回调
        if self.callback:
            self.callback(self.processed_count, self.total_elements)

        return batch_data, start_idx, end_idx

    def reset(self) -> None:
        """重置处理器状态"""
        super().reset()
        self.processed_count = 0


if __name__ == '__main__':
    # 创建列表批处理器
    processor = ListBatchProcessor(range(100), batch_size=15)

    # 迭代处理批次
    for batch, start, end in processor:
        print(f"批次 {processor.current_batch}: {start}-{end}, 大小: {len(batch)}")
        if processor.current_batch == 1:
            print(f"第一批次数据: {batch[:5]}...")

    # 创建测试数据
    data = [
        {"id": i, "category": "A" if i % 2 == 0 else "B", "value": i * 10}
        for i in range(100)
    ]

    # 创建字典批处理器 - 修正参数
    processor = DictBatchProcessor(
        data,
        batch_size=10,
        excls_filter="value",  # 使用数值字段
    )

    # 迭代处理批次
    for batch, start, end in processor:
        print(f'数据：{batch}')
        print(f"批次 {processor.current_batch}: {start}-{end}, 大小: {len(batch)}")
        print(f"第一批次类别分布: {set(item['category'] for item in batch)}")
